{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to a brute-force modelling approach of Pair Distribution Function Data\n",
    "\n",
    "This script first import modules and defines functions.\n",
    "\n",
    "Afterwards, a parameter range can be set for the structures that is under investigation. In this case we investigate structures up to 200 atoms with 7 different structure types: \n",
    "\n",
    "Cubic (sc), body-centered cubic (bcc), face-centered cubic (fcc), hexagonal closed packed (hcp), decahedral, icosahedral and octahedral \n",
    "\n",
    "## Import modules and define functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize.minpack import leastsq\n",
    "from multiprocessing import Pool\n",
    "from diffpy.Structure import Structure, Atom\n",
    "from diffpy.srfit.pdf import PDFContribution, DebyePDFGenerator, PDFParser\n",
    "from diffpy.srfit.fitbase import Profile, FitRecipe, FitContribution, FitResults\n",
    "from ase.cluster.decahedron import Decahedron\n",
    "from ase.cluster.icosahedron import Icosahedron\n",
    "from ase.cluster.octahedron import Octahedron\n",
    "from ase.cluster.cubic import FaceCenteredCubic, BodyCenteredCubic, SimpleCubic\n",
    "from ase.lattice.hexagonal import HexagonalClosedPacked\n",
    "from carbontracker.tracker import CarbonTracker\n",
    "\n",
    "def fit(Qmin, Qmax, Qdamp, cluster, PDFfile, plot):\n",
    "    \"\"\"Function that Fits an experimental PDF using a specific cluster. Set plot=True to get fit parameters.\"\"\"\n",
    "    # Create a PDF contribution as before\n",
    "    pdfprofile = Profile()\n",
    "    pdfparser = PDFParser()\n",
    "    pdfparser.parseFile(PDFfile)\n",
    "    pdfprofile.loadParsedData(pdfparser)\n",
    "    pdfprofile.setCalculationRange(xmin = 1.5, xmax = 20, dx=0.1)\n",
    "\n",
    "    # Setup the PDFgenerator that calculates the PDF from the model\n",
    "    #Generator for first cluster\n",
    "    pdfgenerator = DebyePDFGenerator(\"G\")\n",
    "    pdfgenerator.setQmax(Qmax)\n",
    "    pdfgenerator.setQmin(Qmin)\n",
    "    pdfgenerator._calc.evaluatortype = 'OPTIMIZED'\n",
    "\n",
    "    # Input the data files\n",
    "    pdfgenerator.setStructure(cluster, periodic=False)\n",
    "\n",
    "    # Add the profile and generator the the PDFcontribution\n",
    "    pdfcontribution = FitContribution(\"pdf\")\n",
    "    pdfcontribution.setProfile(pdfprofile, xname=\"r\")\n",
    "    pdfcontribution.addProfileGenerator(pdfgenerator)\n",
    "    pdfcontribution.setEquation(\"scale*G\")\n",
    "\n",
    "    # Moving on\n",
    "    recipe = FitRecipe()\n",
    "    recipe.addContribution(pdfcontribution)\n",
    "    recipe.addVar(pdfcontribution.scale, 0.1, tag = \"scale\")\n",
    "    recipe.restrain(\"scale\", lb=0, ub=1e99, sig=0.001)\n",
    "    pdfgenerator.qdamp.value = Qdamp\n",
    "\n",
    "    # Add ADP for the cluster\n",
    "    phase_struc = pdfgenerator.phase\n",
    "    atoms = phase_struc.getScatterers()\n",
    "\n",
    "    # Make latice to the phase\n",
    "    lat = phase_struc.getLattice()\n",
    "    # Make new variable zoomscale\n",
    "    recipe.newVar(\"zoomscale\", 1.00, tag = \"lat\")\n",
    "    recipe.constrain(lat.a, 'zoomscale')\n",
    "    recipe.constrain(lat.b, 'zoomscale')\n",
    "    recipe.constrain(lat.c, 'zoomscale')\n",
    "\n",
    "    # We create the variables of ADP and assign the initial value to them. In this\n",
    "    # example, we use isotropic ADP for all atoms\n",
    "    Biso = recipe.newVar(\"Biso\", value=1, tag = 'ADP')\n",
    "    recipe.restrain(Biso, lb=0, ub=3, sig = 0.01)\n",
    " \n",
    "    # For all atoms in the structure model, we constrain their Biso according to their species \n",
    "    for atom in atoms:\n",
    "        if atom.element == cluster.element[0]:\n",
    "            recipe.constrain(atom.Biso, Biso)\n",
    "        \n",
    "    recipe.clearFitHooks()\n",
    "    #     Tune PDF\n",
    "    recipe.fix(\"all\")\n",
    "    recipe.free(\"scale\")\n",
    "    leastsq(recipe.residual, recipe.getValues())   \n",
    "    recipe.free(\"ADP\")\n",
    "    leastsq(recipe.residual, recipe.getValues())   \n",
    "    recipe.free(\"lat\")\n",
    "    leastsq(recipe.residual, recipe.getValues())   \n",
    "\n",
    "    # We calculate the goodness-of-fit, Rwp\n",
    "    g = recipe.pdf.profile.y\n",
    "    gcalc = recipe.pdf.evaluate()\n",
    "    rfactor = np.sqrt(sum((g - gcalc)**2) / sum((g)**2))\n",
    "    \n",
    "    if plot:\n",
    "        print (\"FIT RESULTS\\n\")\n",
    "        res = FitResults(recipe)\n",
    "        print (res)\n",
    "\n",
    "        # Plot the observed and refined PDF.\n",
    "        # Get the experimental data from the recipe\n",
    "        r = recipe.pdf.profile.x\n",
    "        gobs = recipe.pdf.profile.y\n",
    "\n",
    "        # Get the calculated PDF and compute the difference between the calculated and measured PDF\n",
    "        gcalc = recipe.pdf.evaluate()\n",
    "        baseline = 1.1 * gobs.min()\n",
    "        gdiff = gobs - gcalc\n",
    "\n",
    "        # Plot!\n",
    "        plt.figure()\n",
    "        plt.plot(r, gobs, 'bo', label=\"G(r) data\")\n",
    "        plt.plot(r, gcalc, 'r-', label=\"G(r) fit\")\n",
    "        plt.plot(r, gdiff + baseline, 'g-', label=\"G(r) diff\")\n",
    "        plt.plot(r, np.zeros_like(r) + baseline, 'k:')\n",
    "        plt.xlabel(r\"$r (\\AA)$\")\n",
    "        plt.ylabel(r\"$G (\\AA^{-2})$\")\n",
    "        plt.legend()\n",
    "\n",
    "        plt.show()\n",
    "    return rfactor\n",
    "  \n",
    "def exp_file_parameters(file):\n",
    "    \"\"\"Function that takes an gr file as input in PDFGui format and extract:\n",
    "    Qmin, Qmax, atom, lattice constant and filename.\"\"\"\n",
    "    f = open(file, \"r\")\n",
    "    file_text = f.read()\n",
    "    Qmin = float(file_text[file_text.find(\"qmin\")+7:file_text.find(\"qmin\")+10])\n",
    "    Qmax = float(file_text[file_text.find(\"qmax \")+7:file_text.find(\"qmax \")+11])\n",
    "    atom = file_text[file_text.find(\"composition\")+14:file_text.find(\"composition\")+16]\n",
    "    try:\n",
    "        data = np.loadtxt(file)\n",
    "    except ValueError:\n",
    "        data = np.loadtxt(file, skiprows=24)\n",
    "    lc = data[data[:,1].argmax(),0]\n",
    "    return Qmin, Qmax, atom, lc, file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use the brute-force approach on all dataset in the \"Data\" folder "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tracker = CarbonTracker(epochs=1, log_dir='./CO2_consumption/')\n",
    "tracker.epoch_start()\n",
    "\n",
    "data_path = \"Data/\"\n",
    "size_threshold = 200\n",
    "cores = 1\n",
    "# Parameter range for SC, BCC and FCC\n",
    "h_min = 1 #Min 1\n",
    "h_max = 10\n",
    "k_min = 1 # Min 1\n",
    "k_max = 10\n",
    "l_min = 1 #Min 1\n",
    "l_max = 10\n",
    "# Parameter range for HCP\n",
    "lc1_min = 1 #Min 1\n",
    "lc1_max = 10\n",
    "lc2_min = 1 # Min 1\n",
    "lc2_max = 10\n",
    "size1_min = 1 #Min 1\n",
    "size1_max = 10\n",
    "size2_min = 1 #Min 1\n",
    "size2_max = 10\n",
    "size3_min = 1 #Min 1\n",
    "size3_max = 10\n",
    "# Parameter range for Icosahedral\n",
    "shell_min = 1\n",
    "shell_max = 10\n",
    "# Parameter range for Decahedral\n",
    "p_start =1 #min 1\n",
    "p_end=10\n",
    "q_start=1 #min 1\n",
    "q_end=10\n",
    "r_start=1\n",
    "r_end=10\n",
    "# Parameter range for Octahedral\n",
    "length_min = 1\n",
    "length_max = 10\n",
    "\n",
    "# Get all the experimental files\n",
    "files = sorted(glob.glob(data_path+\"*.gr\"))\n",
    "files = files[:8]\n",
    "\n",
    "def Brute_Force_Metal_Mining(file):\n",
    "    i = 0\n",
    "    values = {}\n",
    "    print (\"Brute Force Metal Mining on File: \", file)\n",
    "    # Determine the exp file parameters \n",
    "    Qmin, Qmax, atom, lc, PDFFile = exp_file_parameters(file)\n",
    "    # Brute Force metal mining of the file\n",
    "    \n",
    "    print (\"Trying out FCC, BCC and SC\")\n",
    "    # Make FCC, BCC and SC and fit\n",
    "    for h in range(h_min, h_max):\n",
    "        for k in range(k_min, k_max):\n",
    "            for l in range(l_min, l_max):\n",
    "                surfaces = [(1, 0, 0), (1, 1, 0), (1, 1, 1)]\n",
    "                layers = [h, k, l]\n",
    "\n",
    "                # Make FCC and fit\n",
    "                stru1 = FaceCenteredCubic(atom, surfaces, layers, latticeconstant=2*np.sqrt(0.5*lc**2))\n",
    "                xyz1 = stru1.get_positions()\n",
    "                if len(xyz1) < size_threshold:\n",
    "                    Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "                    rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=file, plot=False)\n",
    "                    values[i] = {}\n",
    "                    values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "                    values[i]['structuretype'] = 'FCC'\n",
    "                    values[i]['Rwp'] = rw\n",
    "                    values[i]['Atoms'] = len(xyz1)\n",
    "                    values[i]['h'] = h\n",
    "                    values[i]['k'] = k\n",
    "                    values[i]['l'] = l\n",
    "                    i += 1\n",
    "\n",
    "                # Make BCC and fit\n",
    "                stru1 = BodyCenteredCubic(atom, surfaces, layers, latticeconstant=lc)\n",
    "                xyz1 = stru1.get_positions()\n",
    "                if len(xyz1) < size_threshold:\n",
    "                    Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "                    rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=PDFFile, plot=False)\n",
    "                    values[i] = {}\n",
    "                    values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "                    values[i]['structuretype'] = 'BCC'\n",
    "                    values[i]['Rwp'] = rw\n",
    "                    values[i]['Atoms'] = len(xyz1)\n",
    "                    values[i]['h'] = h\n",
    "                    values[i]['k'] = k\n",
    "                    values[i]['l'] = l\n",
    "                    i += 1\n",
    "\n",
    "                # Make SC and fit\n",
    "                stru1 = SimpleCubic(atom, surfaces, layers, latticeconstant=lc)\n",
    "                xyz1 = stru1.get_positions()\n",
    "                if len(xyz1) < size_threshold:\n",
    "                    Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "                    rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=PDFFile, plot=False)\n",
    "                    values[i] = {}\n",
    "                    values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "                    values[i]['structuretype'] = 'SC'\n",
    "                    values[i]['Rwp'] = rw\n",
    "                    values[i]['Atoms'] = len(xyz1)\n",
    "                    values[i]['h'] = h\n",
    "                    values[i]['k'] = k\n",
    "                    values[i]['l'] = l\n",
    "                    i += 1\n",
    "\n",
    "    print (\"Trying out HCP\")\n",
    "    # Make HCP and Fit\n",
    "    for lc1 in range(lc1_min, lc1_max):\n",
    "        for lc2 in range(lc2_min, lc2_max):\n",
    "            for size1 in range(size2_min, size1_max):\n",
    "                for size2 in range(size2_min, size2_max):\n",
    "                    for size3 in range(size3_min, size3_max):\n",
    "                        layers = [h, k, l]\n",
    "                        stru1 = HexagonalClosedPacked(atom, latticeconstant=(lc1, lc2), size=(size1, size2, size3))\n",
    "                        xyz1 = stru1.get_positions()\n",
    "                        if len(xyz1) < size_threshold:\n",
    "                            Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "                            rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=PDFFile, plot=False)\n",
    "                            values[i] = {}\n",
    "                            values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "                            values[i]['structuretype'] = 'HCP'\n",
    "                            values[i]['Rwp'] = rw\n",
    "                            values[i]['Atoms'] = len(xyz1)\n",
    "                            values[i]['lc1'] = lc1\n",
    "                            values[i]['lc2'] = lc2\n",
    "                            values[i]['size1'] = size1\n",
    "                            values[i]['size2'] = size2\n",
    "                            values[i]['size3'] = size3\n",
    "                            i += 1\n",
    "\n",
    "    print (\"Trying out Icosahedral\")\n",
    "    # Make Icosahedral and Fit\n",
    "    for shell in range(shell_min,shell_max):\n",
    "        stru1 = Icosahedron(atom, shell,latticeconstant=lc)\n",
    "        xyz1 = stru1.get_positions()\n",
    "        if len(xyz1) < size_threshold:\n",
    "            Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "            rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=PDFFile, plot=False)\n",
    "            values[i] = {}\n",
    "            values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "            values[i]['structuretype'] = 'Icosahedral'\n",
    "            values[i]['Rwp'] = rw\n",
    "            values[i]['Atoms'] = len(xyz1)\n",
    "            values[i]['shell'] = shell\n",
    "            i += 1\n",
    "\n",
    "    print (\"Trying out Decahedral\")  \n",
    "    # Make Decahedral and Fit  \n",
    "    for p in range(p_start,p_end):\n",
    "        for q in range(q_start,q_end):\n",
    "            for r in range(r_start, r_end):\n",
    "                stru1 = Decahedron(atom, p, q, r,latticeconstant=lc)\n",
    "                xyz1 = stru1.get_positions()\n",
    "                if len(xyz1) < size_threshold:\n",
    "                    Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "                    rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=PDFFile, plot=False)\n",
    "                    values[i] = {}\n",
    "                    values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "                    values[i]['structuretype'] = 'Decahedral'\n",
    "                    values[i]['Rwp'] = rw\n",
    "                    values[i]['Atoms'] = len(xyz1)\n",
    "                    values[i]['p'] = p\n",
    "                    values[i]['q'] = q\n",
    "                    values[i]['r'] = r\n",
    "                    i += 1\n",
    "\n",
    "    print (\"Trying out Octahedral\")\n",
    "    # Make Octahedral and Fit  \n",
    "    for length in range(length_min, length_max):\n",
    "        stru1 = Octahedron(atom, length, latticeconstant=2*np.sqrt(0.5*lc**2))\n",
    "        xyz1 = stru1.get_positions()\n",
    "        if len(xyz1) < size_threshold:\n",
    "            Bi1 = Structure([Atom(atom, xi) for xi in xyz1])\n",
    "            rw = fit(Qmin = Qmin, Qmax = Qmax, Qdamp = 0.03, cluster = Bi1, PDFfile=PDFFile, plot=False)\n",
    "            values[i] = {}\n",
    "            values[i]['dataset'] = file.replace(data_path, \"\")[:-3]\n",
    "            values[i]['structuretype'] = 'Octahedral'\n",
    "            values[i]['Rwp'] = rw\n",
    "            values[i]['Atoms'] = len(xyz1)\n",
    "            values[i]['p'] = length\n",
    "            i += 1\n",
    "\n",
    "\n",
    "    df = pd.DataFrame(values)\n",
    "    df.to_csv(save_path+ file.replace(data_path, \"\")[:-3] + \".csv\")\n",
    "    return None\n",
    "\n",
    "p = Pool(processes=cores)\n",
    "results = p.map(Brute_Force_Metal_Mining, files)\n",
    "p.close()\n",
    "p.join()\n",
    "results\n",
    "\n",
    "tracker.epoch_end()\n",
    "tracker.stop()\n",
    "parser.print_aggregate(log_dir=\"./CO2_consumption/\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
